import pickle
from cpython.unicode cimport PyUnicode_EncodeUTF16, PyUnicode_DecodeUTF16, PyUnicode_EncodeUTF8, PyUnicode_DecodeUTF8
include "uint_hash_table.pyx"
#-----------------------------------------------------------------------------------------------------------------------
cpdef resize_array(array arr, u64 cur, u64 extend_size):
    cdef u64 larr = Py_SIZE(arr)
    resize_smart(arr, cur+extend_size)
    return arr

cpdef u64 encode_u64(u64 num, array arr, u64 cur):
    cdef u64 bytecount, larr = Py_SIZE(arr)
    #cdef u8[9] encoded_num
    cdef u8 *ptr = arr.data.as_uchars + cur, *end = ptr + larr
    if ptr + 9 < end:
        pass
    else:
        arr = resize_array(arr, cur, 1<<20)
        ptr = arr.data.as_uchars + cur
    bytecount = _encode_num(num, ptr)
    return cur+bytecount


#-----------------------------------------------------------------------------------------------------------------------

#-----------------------------------------------------------------------------------------------------------------------

cpdef u64 encode_uint_set(UintSet ids, array arr, u64 cur):
    cdef u64 max_v = ids._pop_max_value()
    cur=encode_u64(ids.count, arr, cur)
    cur=encode_u64(max_v, arr, cur)
    cdef u64 i=0, v
    while(v!=0):
        v = ids._next(&i)
        cur = encode_u64(v, arr, cur)
    return cur

cdef class RefIdsABC:
    cdef readonly str name

cdef class RefIds(RefIdsABC):
    cdef readonly UintSet ids
    cdef RefId get_item_(self, u64 i):
       return RefId().set(self.name, self._get_item(i))

cdef class RefIdsMap(RefIdsABC):
    cdef readonly UintMap map

cdef class RefId(RefIdsABC):
    cdef readonly u64 id
    cdef RefId set(self, str name, u64 id):
        self.name = name
        self.id = id
        return self

cpdef u64 encode_ref_ids_map(RefIdsMap ids, array arr, u64 cur, dict name_map_num_encode):
    cdef u64 i = 0, key = ids.map._next(&i)
    while (key != 0):
        value=ids.map.values(i)
        encode_u64(key, arr, cur)
        encode_item_or_items(value, arr, cur, name_map_num_encode)
    return cur
cpdef u64 encode_ref_ids(RefIds ids, array arr, u64 cur, dict name_map_num_encode):
    cdef U64 type_num = name_map_num_encode[id.name]
    cdef u64 cur = encode_u64(type_num, arr, cur)
    return encode_uint_set(ids.ids, arr, cur)


cpdef u64 encode_ref_id(RefId id, array arr, u64 cur, dict name_map_num_encode):
    cdef U64 type_num = name_map_num_encode[id.name]
    cdef u64 cur = encode_u64(type_num, arr, cur)
    return encode_u64((id.id<<1), arr, cur)

cdef class RefDecoder:
    cdef readonly str name
    cdef readonly Buffer buf
    def __call__(self, array arr, u64 cur):
        return self.decode(arr, cur)
    cpdef decode(self, array arr, u64 cur):
        cdef ptr = arr.data.as_uchars+cur
        cdef u64 num = decode_u64(arr, cur, self.buf), count, id, max_id
        cur = self.buf.size
        cdef u32* u32s
        cdef u64* u64s
        if num&1:
            count = num>>1
            if count>0:
                max_id = decode_u64(arr, cur, self.buf)
                cur = self.buf.size
                if max_id<=UINT32_MAX:
                    u32s = <u32*>py_malloc(count*4)
                    u32s[0]=<u32>max_id
                    for i in range(1,count):
                        id = decode_u64(arr, cur, self.buf)
                        u32s[i] = <u32>id
                    return RefIds().set(u32s, 2, count)
                else:
                    u64s = <u64*>py_malloc(count*8)
                    u64s[0]=max_id
                    for i in range(1,count):
                        id = decode_u64(arr, cur, self.buf)
                        u64s[i] = id
                    return RefIds().set(u64s, 3, count)
        else:
            id = num>>1
            return RefId().set(self.name, id)

#-----------------------------------------------------------------------------------------------------------------------

cdef bytes ONE_NUM=encode_num_to_bytes(1), ZERO_NUM=encode_num_to_bytes(0)
cdef U64 NONE=get_U64(0),OBJ=get_U64(1), TUPLE=get_U64(2), LIST=get_U64(3), SET=get_U64(4), FROZENSET=get_U64(5)
cdef U64 DICT=get_U64(6), Ids=get_U64(7), RefId=get_U64(8), RefIds=get_U64(9)
cpdef u64 encode_seq_items(items, array arr, u64 cur, dict type_map_num_encode):
    for item in items:
        cur = encode_item_or_items(item, arr, cur, type_map_num_encode)
    return cur

cpdef u64 encode_map_items(items, array arr, u64 cur, dict type_map_num_encode):
    for key, value in items:
        cur = encode_item_or_items(key, arr, cur, type_map_num_encode)
        cur = encode_item_or_items(value, arr, cur, type_map_num_encode)
    return cur

cpdef u64 encode_item(item, type t, array arr, u64 cur, dict type_map_num_encode):
    cdef const u8[:] value
    cdef U64 num
    cdef u64 count, len
    if t is RefIds:
        return encode_ref_ids_map(item, arr, cur, type_map_num_encode)
    elif t is RefId:
        return encode_ref_id(item, arr, cur, type_map_num_encode)
    elif t is RefIdsMap:
        return encode_ref_ids_map(item, arr, cur, type_map_num_encode)
    #
    try:
        num, encode = type_map_num_encode[t]
    except KeyError:
        type_num = ONE_NUM
        encode = pickle.dumps
    #
    cur = encode_u64(type_num.val, arr, cur)
    try:
        value = encode(item)
    except Exception as e:
        raise RuntimeError from e
    encode_u64(value.nbytes, arr, cur)
    #
    return encode_buffer( &value[0], value.nbytes, arr, cur)

cpdef u64 encode_item_or_items(item_or_items, array arr, u64 cur, dict type_map_num_encode):
    cdef u64 cur
    cdef list l
    cdef tuple tp
    cdef set s
    cdef frozenset fs
    cdef dict d
    if item_or_items is None:
        cur = encode_u64(0, arr, cur)
        return cur
    #
    t = type(item_or_items)
    if t is tuple:
        tp = item_or_items
        cur = encode_u64(1, arr, cur)
        cur = encode_u64(len(tp), arr, cur)
        return encode_seq_items(item_or_items, arr, cur, type_map_num_encode)
    elif t is list :
        l = item_or_items
        cur = encode_u64(2, arr, cur)
        cur = encode_u64(len(l), arr, cur)
        return encode_seq_items(item_or_items, arr, cur, type_map_num_encode)
    elif t is set :
        s = item_or_items
        cur = encode_u64(3, arr, cur)
        cur = encode_u64(len(s), arr, cur)
        return encode_seq_items(item_or_items, arr, cur, type_map_num_encode)
    elif t is frozenset:
        fs = item_or_items
        cur = encode_u64(4, arr, cur)
        cur = encode_u64(len(fs), arr, cur)
        return encode_seq_items(item_or_items, arr, cur, type_map_num_encode)
    elif t is dict:
        d=item_or_items
        cur = encode_u64(5, arr, cur)
        cur = encode_u64(len(d), arr, cur)
        return encode_map_items(item_or_items, arr, cur, type_map_num_encode)
    else:
        return encode_item(item_or_items, t, arr, cur, type_map_num_encode)


#-----------------------------------------------------------------------------------------------------------------------
cdef u64 encode_buffer(void* data, u64 size, array arr, u64 cur):
    cdef u64 larr = Py_SIZE(arr), lvalue=size
    if cur+lvalue<larr:
        pass
    else:
        arr = resize_array(arr, cur, 1 << 20)
    cdef u8 *ptr = arr.data.as_uchars + cur
    memcpy(ptr, data, lvalue)
    return cur+lvalue

cpdef u64 encode_bytes(bytes value, array arr, u64 cur):
    cdef u64 larr = Py_SIZE(arr), lvalue = len(value)
    if cur+lvalue<larr:
        pass
    else:
        arr = resize_array(arr, cur, 1 << 20)
    cdef u8 *ptr = arr.data.as_uchars + cur
    memcpy(ptr, <char*>value, lvalue)
    return cur+lvalue
#-----------------------------------------------------------------------------------------------------------------------

cpdef u64 encode_u64_key(key, array arr, u64 cur, dict type_map_num_encode):
    cdef U64 k=key
    return encode_u64(k.val, arr, cur)


#-----------------------------------------------------------------------------------------------------------------------
cpdef u64 encode_fix_fields(fix_fields, array arr, u64 cur, dict type_map_num_encode):
    for value in fix_fields:
        cur = encode_item_or_items(value, arr, cur, type_map_num_encode)
    return cur

cdef u64 encode_var_fields(dict var_fields, array arr, u64 cur, dict type_map_num_encode, dict field_map_num):
    cdef U64 field_num
    if var_fields:
        cur = encode_u64(len(var_fields), arr, cur)
        for field, value in var_fields:
            field_num = field_map_num[field]
            cur = encode_u64(field_num.val, arr, cur)
    else:
        cur = encode_u64(ZERO_NUM.val , arr, cur)
    return cur
#-----------------------------------------------------------------------------------------------------------------------
cpdef encode_obj_kvs(kvs, u64 blocksize, bint key_type_u64, dict type_map_num_encode, u64 init_array_len=1<<20):
    cdef u64  cur=0
    cdef array arr=new_array(init_array_len), foot_arr
    cdef list blocks=[], split_keys=[]
    cdef tuple fix_fields
    cdef dict var_fields
    cdef str field
    for key, fix_fields, var_fields in kvs:
        if not key_type_u64:
            cur = encode_item(key, arr, cur, type_map_num_encode)
        else:
            cur = encode_u64_key(key, arr, cur, type_map_num_encode)
        cur = encode_fix_fields(fix_fields, arr, cur, type_map_num_encode)
        cur = encode_var_fields(var_fields, arr, cur, type_map_num_encode)
        if cur<blocksize:
            continue
        else:
            resize(arr, cur)
            blocks.append(arr)
            split_keys.append(key)
            #
            arr = new_array(init_array_len)
            cur=0
    #
    if 0<cur<blocksize:
        if blocks:
            foot_arr = blocks[len(blocks)-1]
            array_extend_buffer(foot_arr, arr.data.as_uchars, cur)
            resize(arr,0)
            arr = None
            split_keys[len(split_keys)-1] = key
    elif cur<blocksize:
        resize(arr, cur)
        blocks.append(arr)
        split_keys.append(key)
        arr = None
    else:
        arr=None
    return blocks, split_keys, arr
#-----------------------------------------------------------------------------------------------------------------------
cpdef u64 decode_u64(array data, u64 cur, Buffer buf):
    if cur<Py_SIZE(data):
        pass
    else:
        raise BufferError
    #
    cdef u64 bytecount, v= _decode_num(data.data.as_uchars+cur, &bytecount)
    buf.size = cur+bytecount
    return v

cpdef decode_u64_key(array data, u64 cur, Buffer buf):
    return get_U64(decode_u64(data, cur, buf))

cpdef decode_obj_item(array data, u64 cur, Buffer buf, tuple decode_funcs):
    if cur<Py_SIZE(data):
        pass
    else:
        raise BufferError
    #
    cdef u8* ptr =data.data.as_uchars
    cdef u64 bytecount, type_num = _decode_num(ptr+cur, &bytecount)
    cur+=bytecount
    cdef u64 len=_decode_num(ptr+cur, &bytecount)
    cur+=bytecount
    buf.set(ptr+cur, len)
    decode = decode_funcs[type_num]
    key = decode(buf)
    buf.size = cur+len
#-----------------------------------------------------------------------------------------------------------------------
cpdef dict decode_obj_kvs(array data, tuple decode_funcs, tuple fix_fields, tuple var_fields, bint key_type_u64 ):
    cdef u64 cur = 0, l=Py_SIZE(data), bytecount, var_count, var_i, var_num
    cdef Buffer buf = Buffer()
    cdef dict kv, kvs={}
    while( cur<l ):
        kv={}
        if cur<l:pass
        else:raise BufferError
        #
        if key_type_u64:
            key = decode_u64_key(data, cur, buf)
        else:
            key = decode_obj_item(data, cur, buf, decode_funcs)
        cur = buf.size
        #
        for field in fix_fields:
            if cur < l:
                pass
            else:
                raise BufferError
            value = decode_obj_item(data, cur, buf, decode_funcs)
            cur = buf.size
            kv[field] = value
        #
        var_count = decode_u64(data, cur, buf)
        cur = buf.size
        var_i=0
        while(var_i<var_count):
            var_num = decode_u64(data, cur, buf)
            field = var_fields[var_num]
            cur = buf.size
            #
            value = decode_obj_item(data, cur, buf, decode_funcs)
            cur = buf.size
            #
            var_i += 1
            kv[field] = value
        kvs[key] = kv
    assert cur==l
    return kvs







#-----------------------------------------------------------------------------------------------------------------------






